package com.ruoyi.model.controller;

import com.ruoyi.common.annotation.Log;
import com.ruoyi.common.core.controller.BaseController;
import com.ruoyi.common.core.domain.AjaxResult;
import com.ruoyi.common.core.page.TableDataInfo;
import com.ruoyi.common.enums.BusinessType;
import com.ruoyi.common.utils.poi.ExcelUtil;
import com.ruoyi.model.domain.TrainTask;
import com.ruoyi.model.service.ITrainTaskService;
import org.apache.commons.io.FileUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.*;

import javax.servlet.http.HttpServletResponse;
import java.io.File;
import java.io.IOException;
import java.net.URLEncoder;
import java.util.List;

/**
 * 模型训练任务Controller
 * 
 * @author ruoyi
 * @date 2025-03-10
 */
@RestController
@RequestMapping("/task/task")
public class TrainTaskController extends BaseController {
    @Autowired
    private ITrainTaskService trainTaskService;

    @Value("${ruoyi.profile}")
    private String uploadPath;

    /**
     * 查询模型训练任务列表
     */
    @PreAuthorize("@ss.hasPermi('task:task:list')")
    @GetMapping("/list")
    public TableDataInfo list(TrainTask trainTask) {
        startPage();
        List<TrainTask> list = trainTaskService.selectTrainTaskList(trainTask);
        return getDataTable(list);
    }

    /**
     * 导出模型训练任务列表
     */
    @PreAuthorize("@ss.hasPermi('task:task:export')")
    @Log(title = "模型训练任务", businessType = BusinessType.EXPORT)
    @PostMapping("/export")
    public void export(HttpServletResponse response, TrainTask trainTask) {
        List<TrainTask> list = trainTaskService.selectTrainTaskList(trainTask);
        ExcelUtil<TrainTask> util = new ExcelUtil<TrainTask>(TrainTask.class);
        util.exportExcel(response, list, "模型训练任务数据");
    }

    /**
     * 获取模型训练任务详细信息
     */
    @PreAuthorize("@ss.hasPermi('task:task:query')")
    @GetMapping(value = "/{id}")
    public AjaxResult getInfo(@PathVariable("id") Long id) {
        return success(trainTaskService.selectTrainTaskById(id));
    }

    /**
     * 新增模型训练任务
     */
    @PreAuthorize("@ss.hasPermi('task:task:add')")
    @Log(title = "模型训练任务", businessType = BusinessType.INSERT)
    @PostMapping
    public AjaxResult add(@RequestBody TrainTask trainTask) {
        return toAjax(trainTaskService.insertTrainTask(trainTask));
    }

    /**
     * 修改模型训练任务
     */
    @PreAuthorize("@ss.hasPermi('task:task:edit')")
    @Log(title = "模型训练任务", businessType = BusinessType.UPDATE)
    @PutMapping
    public AjaxResult edit(@RequestBody TrainTask trainTask) {
        return toAjax(trainTaskService.updateTrainTask(trainTask));
    }

    /**
     * 删除模型训练任务
     */
    @PreAuthorize("@ss.hasPermi('task:task:remove')")
    @Log(title = "模型训练任务", businessType = BusinessType.DELETE)
    @DeleteMapping("/{ids}")
    public AjaxResult remove(@PathVariable Long[] ids) {
        return toAjax(trainTaskService.deleteTrainTaskByIds(ids));
    }

    /**
     * 开始训练任务
     */
    @PreAuthorize("@ss.hasPermi('task:task:start')")
    @Log(title = "模型训练任务", businessType = BusinessType.UPDATE)
    @PutMapping("/start/{id}")
    public AjaxResult startTraining(@PathVariable("id") Long id) {
        return toAjax(trainTaskService.startTraining(id));
    }

    /**
     * 停止训练任务
     */
    @PreAuthorize("@ss.hasPermi('task:task:stop')")
    @Log(title = "模型训练任务", businessType = BusinessType.UPDATE)
    @PutMapping("/stop/{id}")
    public AjaxResult stopTraining(@PathVariable("id") Long id) {
        return toAjax(trainTaskService.stopTraining(id));
    }

    /**
     * 下载训练结果文件
     */
    @PreAuthorize("@ss.hasPermi('task:task:download')")
    @GetMapping("/download/{id}")
    public void download(@PathVariable("id") Long id, HttpServletResponse response) throws IOException {
        TrainTask trainTask = trainTaskService.selectTrainTaskById(id);
        if (trainTask == null || !StringUtils.hasText(trainTask.getResultFile())) {
            throw new RuntimeException("未找到可下载的文件");
        }

        String filePath = uploadPath + "/" + trainTask.getResultFile();
        System.out.println(filePath);
        File file = new File(filePath);
        if (!file.exists() || !file.isFile()) {
            throw new RuntimeException("文件不存在或已损坏");
        }

        try {
            response.reset();
            response.setContentType("application/x-zip-compressed");
            response.setCharacterEncoding("utf-8");
            response.setHeader("Content-Disposition", "attachment;filename=" +
                    URLEncoder.encode(trainTask.getName() + "_result.zip", "UTF-8"));
            response.setHeader("Content-Length", String.valueOf(file.length()));

            FileUtils.copyFile(file, response.getOutputStream());
            response.getOutputStream().flush();
        } catch (IOException e) {
            throw new RuntimeException("文件下载失败", e);
        }
    }
}
